package com.cn.lp.ai.factory;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.ResponseEntity;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;

public class WenxinChatModel {
    private static final Logger logger = LoggerFactory.getLogger(WenxinChatModel.class);
    public final RetryTemplate retryTemplate;
    private final WenxinApi qianFanApi;

    public WenxinChatModel(WenxinApi qianFanApi) {
        this(qianFanApi, RetryUtils.DEFAULT_RETRY_TEMPLATE);
    }

    public WenxinChatModel(WenxinApi qianFanApi, RetryTemplate retryTemplate) {
        Assert.notNull(qianFanApi, "CustomerQianFanApi must not be null");
        Assert.notNull(retryTemplate, "RetryTemplate must not be null");
        this.qianFanApi = qianFanApi;
        this.retryTemplate = retryTemplate;
    }

    public CustomerQianFanApiParams.TextToImageResponse textToImage(String text) {
        CustomerQianFanApiParams.TextToImageRequest request = this.createTextToImageRequest(text);
        return this.retryTemplate.execute((ctx) -> {
            ResponseEntity<CustomerQianFanApiParams.TextToImageResponse> completionEntity = this.qianFanApi.textToImage(request);
            CustomerQianFanApiParams.TextToImageResponse chatCompletion = completionEntity.getBody();
            if (chatCompletion == null) {
                logger.warn("No data returned for textToImage: {}", text);
                return new CustomerQianFanApiParams.TextToImageResponse(null, "");
            } else {
                return chatCompletion;
            }
        });
    }

    public CustomerQianFanApiParams.TextToImageRequest createTextToImageRequest(String text) {
        return new CustomerQianFanApiParams.TextToImageRequest(text);
    }

}
